from datasets import load_dataset
from transformers import BertTokenizer, BertForSequenceClassification
from transformers import TrainingArguments, Trainer

# 加载 IMDb 数据集
dataset = load_dataset("imdb")

# 训练集、测试集、验证集
train_data = dataset["train"]
test_data = dataset["test"]
print(train_data[1])

# 选择 BERT 预训练模型
model_name = "bert-base-uncased"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForSequenceClassification.from_pretrained(model_name, num_labels=2)

def preprocess_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)

# 预处理数据
train_data = train_data.map(preprocess_function, batched=True)
test_data = test_data.map(preprocess_function, batched=True)

# 选择训练字段
train_data.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
test_data.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])

# 训练参数
training_args = TrainingArguments(
    output_dir="E:/temp/train/results",      # 训练结果保存路径
    evaluation_strategy="epoch", # 每个 epoch 进行一次评估
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,          # 训练轮数
    weight_decay=0.01,           # L2正则化
    logging_dir="E:/temp/train/logs",        # 日志目录
    save_total_limit=2,          # 最多保留2个检查点
)

# Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    eval_dataset=test_data,
)

# 开始训练
trainer.train()

